import torch.nn as nn
import torch


class Loss(nn.Module):
    def __init__(self):
        super(Loss, self).__init__()

    @staticmethod
    def forward(winner, self_play_winner, probas, self_play_probas):
        value_error = (self_play_winner - winner) ** 2
        policy_error = torch.sum((-self_play_probas * (1e-6 + probas).log()), 1)
        total_error = (value_error.view(-1) + policy_error).mean()
        return total_error
